from typing import TypedDict
import uuid
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.types import interrupt, Command

class State(TypedDict):
    text_1: str
    text_2: str

def human_node_1(state: State):
    value = interrupt({"text_to_revise": state["text_1"]})
    return {"text_1": value}


def human_node_2(state: State):
    value = interrupt({"text_to_revise": state["text_2"]})
    return {"text_2": value}


graph_builder = StateGraph(State)
graph_builder.add_node("human_node_1", human_node_1)
graph_builder.add_node("human_node_2", human_node_2)

# Add both nodes in parallel from START
graph_builder.add_edge(START, "human_node_1")
graph_builder.add_edge(START, "human_node_2")

checkpointer = InMemorySaver()
graph = graph_builder.compile(checkpointer=checkpointer)

thread_id = str(uuid.uuid4())
thread_config: RunnableConfig = {"configurable": {"thread_id": thread_id}}

result = graph.invoke(
    {"text_1": "original text 1", "text_2": "original text 2"}, config=thread_config
)
print(result)
print("=================================")
# Resume with mapping of interrupt IDs to values
resume_map = {
    i.interrupt_id: f"human input for prompt {i.value['text_to_revise']}"
    for i in graph.get_state(thread_config).interrupts
}
output = graph.invoke(Command(resume=resume_map), config=thread_config)

print(output)
# > {'text_1': 'edited text for original text 1', 'text_2': 'edited text for original text 2'}